package net.javlov.example.landmarks;

import java.awt.Point;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

import javax.swing.Timer;

import net.javlov.Action;
import net.javlov.Agent;
import net.javlov.EpisodicRewardStepStatistic;
import net.javlov.Option;
import net.javlov.Policy;
import net.javlov.QLearningAgent;
import net.javlov.SarsaAgent;
import net.javlov.Simulator;
import net.javlov.TabularQFunction;
import net.javlov.TabularValueFunction;
import net.javlov.example.ExperimentGUI;
import net.javlov.example.GridLimitedOptionsWorld;
import net.javlov.example.rooms.ReachHallOption;
import net.javlov.policy.EGreedyPolicy;
import net.javlov.world.AgentBody;
import net.javlov.world.Body;
import net.javlov.world.GoalBody;
import net.javlov.world.grid.GridGPSSensor;
import net.javlov.world.grid.GridMove;
import net.javlov.world.grid.GridRewardFunction;
import net.javlov.world.phys2d.Phys2DAgentBody;
import net.javlov.world.phys2d.Phys2DBody;
import net.javlov.world.ui.GridWorldView;
import net.phys2d.raw.shapes.Circle;
import net.phys2d.raw.shapes.StaticBox;

public class Main implements Runnable {
	
	GridLimitedOptionsWorld world;
	GridRewardFunction rf;
	Simulator sim;
	TabularQFunction qf;
	SarsaAgent agent;
	boolean gui;
	int cellwidth, cellheight;
	
	public static void main(String[] args) {
		Main m = new Main();
		m.gui = false;
		m.init();
		m.start();
	}
	
	public void init() {
		cellwidth = 50; cellheight = cellwidth;

		makeWorld();		
		
		List<? extends Option> optionPool = makeOptions();
		world.setOptionPool(optionPool);
		
		agent = makeAgent(optionPool);
		AgentBody aBody = makeAgentBody();
		
		sim = new Simulator();
		sim.setEnvironment(world);
		
		world.add(agent, aBody);
		sim.addAgent(agent);
	}
	
	protected SarsaAgent makeAgent(List<? extends Option> optionPool) {
		qf = TabularQFunction.getInstance(optionPool.size());
		SarsaAgent a = new QLearningAgent(qf, 1);
		a.setLearnRate(0.1);
		Policy pi = new EGreedyPolicy(qf, 0.05, optionPool);
		a.setPolicy(pi);
		a.setSMDPMode(true);
		a.setLearnStateValueFunction(true);
		return a;
	}
	
	protected AgentBody makeAgentBody() {
		Phys2DAgentBody aBody = new Phys2DAgentBody(new Circle(20), 0.5f);
		GridGPSSensor gps = new GridGPSSensor(cellwidth, cellheight);
		gps.setBody(aBody);
		aBody.add(gps);
		return aBody;
	}
	
	protected List<? extends Option> makeOptions() {
		List<Action> primitiveActions = new ArrayList<Action>();
		primitiveActions.add(GridMove.getNorthInstance(world));
		primitiveActions.add(GridMove.getNorthEastInstance(world));
		primitiveActions.add(GridMove.getEastInstance(world));
		primitiveActions.add(GridMove.getSouthEastInstance(world));
		primitiveActions.add(GridMove.getSouthInstance(world));
		primitiveActions.add(GridMove.getSouthWestInstance(world));
		primitiveActions.add(GridMove.getWestInstance(world));
		primitiveActions.add(GridMove.getNorthWestInstance(world));
		
		List<Option> optionPool = new ArrayList<Option>();
		Option o = new ReachLandmarkOption("L1", 100, new Point(4,14), primitiveActions);
		o.setID(optionPool.size());
		optionPool.add( o );
		
		o = new ReachLandmarkOption("L2", 8, new Point(4,8), primitiveActions);
		o.setID(optionPool.size());
		optionPool.add( o );
		
		o = new ReachLandmarkOption("L3", 8, new Point(10,8), primitiveActions);
		o.setID(optionPool.size());
		optionPool.add( o );
		
		o = new ReachLandmarkOption("L4", 8, new Point(10,2), primitiveActions);
		o.setID(optionPool.size());
		optionPool.add( o );
		
		return optionPool;
	}
	
	protected void makeWorld() {
		world = new GridLimitedOptionsWorld(15, 18, cellwidth, cellheight);
		
		Body landmark = new Phys2DBody( new StaticBox(cellwidth, cellheight), 10, true );
		landmark.setLocation(4.5*cellwidth, 14.5*cellheight);
		landmark.setType(Body.UNDEFINED);
		world.addFixedBody(landmark);
		
		landmark = new Phys2DBody( new StaticBox(cellwidth, cellheight), 10, true );
		landmark.setLocation(4.5*cellwidth, 8.5*cellheight);
		landmark.setType(Body.UNDEFINED);
		world.addFixedBody(landmark);
		
		landmark = new Phys2DBody( new StaticBox(cellwidth, cellheight), 10, true );
		landmark.setLocation(10.5*cellwidth, 8.5*cellheight);
		landmark.setType(Body.UNDEFINED);
		world.addFixedBody(landmark);
		
		landmark = new Phys2DBody( new StaticBox(cellwidth, cellheight), 10, true );
		landmark.setLocation(10.5*cellwidth, 2.5*cellheight);
		landmark.setType(Body.UNDEFINED);
		world.addFixedBody(landmark);
		
		rf = new GridRewardFunction();
		world.setRewardFunction(rf);
		world.addCollisionListener(rf);
		
		GoalBody goal = new GoalBody(525, 125);
		goal.setReward(0);
		world.addFixedBody(goal);
	}
	
	public void start() {
		if ( gui ) {
			GridWorldView wv = new GridWorldView(world);
			Timer timer = new Timer(1000/24, wv);
			ExperimentGUI g = new ExperimentGUI("Rooms example", wv, sim);
			timer.start();
			new Thread(this).start();
		} else
			run();
	}

	@Override
	public void run() {
		int episodes = 1000;
		EpisodicRewardStepStatistic stat = new EpisodicRewardStepStatistic(episodes);
		sim.addStatistic(stat);
		sim.init();
		//sim.suspend();
		agent.setInterruptOptions(true);
		sim.runEpisodes(episodes);
		System.out.println(agent.getValueFunction());
		
	}

}

